#!/usr/bin/env python
# coding: utf-8

import os, struct, cPickle as pickle
from twisted.protocols import basic
from twisted.internet import stdio, defer, reactor, protocol

class ParallelsProtocol(basic.Int32StringReceiver):
    pid = os.getpid()
    ppid = os.getppid()

    def stringReceived(self, data):
	self.itemReceived(pickle.loads(data))

    def sendItem(self, data):
	self.sendString(pickle.dumps(data))

    def itemReceived(self, item):
	raise NotImplementedError

class ConsumerProtocol(ParallelsProtocol):
    def __init__(self, runReactor=True):
	assert os.environ.get('Twisted-Parallels') == 'True', \
	    'This program is not being run by a Twisted-Parallels producer.'

	stdio.StandardIO(self)
	self.autoStop = runReactor
	if runReactor: reactor.run()

    def connectionLost(self, reason):
	if self.autoStop:
	    try:
		os.close(0)
		os.close(1)
		os.close(2)
		reactor.stop()
	    except: pass
    
    def itemReceived(self, data):
	response = self.get(data)
	if isinstance(response, defer.Deferred):
	    response.addCallback(self.sendItem)
	else:
	    self.sendItem(response)
    
    def get(self, item):
	raise NotImplementedError

    def __repr__(self):
	return '<ConsumerProtocol pid=%d parent=%d>' % (self.pid, self.ppid)

class Parallels(protocol.ProcessProtocol, ParallelsProtocol):
    def __init__(self, childPath, childArgs=[], _queue=None):
	self.busy = False
	self._queue = _queue
	self.deferred = None
	if _queue: _queue.get().addCallback(self._consume)

	# run the child
	cmdname = 'python'
	cmdargs = [cmdname, childPath] + childArgs
	cmdenviron = os.environ.copy()
	cmdenviron['Twisted-Parallels'] = 'True'
	self.proc = reactor.spawnProcess(self, cmdname, cmdargs,
	    cmdenviron, childFDs={0:'w',1:'r',2:2})
	
    def _consume(self, (item, fakedeferred)):
	deferred = self.put(item)
	deferred.chainDeferred(fakedeferred)
    
    def terminate(self):
	self.proc.signalProcess('TERM')

    def connectionLost(self):
	self.terminate()

    def childDataReceived(self, childFD, data):
	if self.busy and childFD == 1: self.dataReceived(data)

    def itemReceived(self, item):
	self.busy = False
	if self.deferred:
	    self.deferred.callback(item)
	    self.deferred = None
	if self._queue: self._queue.get().addCallback(self._consume)

    def put(self, item):
	assert not self.busy, '%s is busy!' % str(self)
	self.busy = True
	self.deferred = defer.Deferred()
	self.sendItem(item)
	return self.deferred
    
    def __repr__(self):
	return '<Parallels pid=%d child=%d>' % (self.ppid, self.proc.pid)

class Pool:
    def __init__(self, consumerPath, consumerArgs=[], size=5):
	self.procs = []
	self.queue = defer.DeferredQueue()

	for n in xrange(size):
	    proc = Parallels(consumerPath, [str(n)]+consumerArgs, _queue=self.queue)
	    self.procs.append(proc)
	
    def shutdown(self):
	for proc in self.procs:
	    proc.terminate()

    def put(self, item):
	fakedeferred = defer.Deferred()
	self.queue.put((item, fakedeferred))
	return fakedeferred
